import retro 
import numpy as np
import cv2
import neat
import pickle
from random import random

# Play this retro game at this level.
env = retro.make('HangOn-Sms')

imgarray = []

def eval_genomes(genomes, config):
    
    for genome_id, genome in genomes:
        ob = env.reset()
        # ac = env.action_space.sample() # action with generic sample

        # 获取环境中的观测空间的大小
        iny, inx, inc = env.observation_space.shape

        # 将观测空间的大小缩小一倍
        inx = int(inx/8)
        iny = int(iny/8)


        # 创建一个神经网络
        net = neat.nn.recurrent.RecurrentNetwork.create(genome, config)
        # nnOutput = net.activate(imgarray)
        fitness_current = 0  # 当前适应度
        scoreTracker = 0  # 分数追踪器
        timeTracker =0  # 时间追踪器
        leftTracker = 0  # 距离追踪器
        speedTracker = 0  # 速度追踪器
        score = 0  # 分数
        left = 0  # 离目标距离
        speed = 0  # 速度
        time = 0  # 时间
        speed_count = 0  # 速度计数
        time_count = 0  # 时间计数
        left_count = 0  # 距离计数
        frame = 0  # 帧

        done = False

        while not done:
            # 渲染环境
            env.render()
            frame += 1

            # 如果随机数小于-0.01，则执行随机动作
            if random()<-0.01:
                nnOutput = env.action_space.sample() # action with generic sample
            else:
                # 将图像转换为灰度图，并flatten
                ob = cv2.resize(ob, (inx, iny))
                ob = cv2.cvtColor(ob, cv2.COLOR_BGR2GRAY)
                ob = np.reshape(ob, (inx, iny))
                imgarray = ob.flatten()
                # 使用神经网络输出动作
                nnOutput = net.activate(imgarray)

            nnOutput[0] = 0 #刹车
            nnOutput[1] 
            nnOutput[2] 
            nnOutput[3] 
            nnOutput[4] # 进档 
            nnOutput[5] # 退档 
            nnOutput[6] # 向左
            nnOutput[7] # 向右    
            nnOutput[8] = 1  #油门

            # 执行环境step函数，传入nnOutput，得到ob, state, done, info四个值
            ob, state, done, info = env.step(nnOutput)        

            # 获取info中的score，left，speed，time四个值
            score = info['score']
            left = info['left']
            speed = info['speed']
            time = info['time']


            # Add to fitness score if mario gains points on his score.
            if score > 0:
                if score > scoreTracker:
                    fitness_current = score * 10
                    scoreTracker = score
            
            # If time > 60 seconds have add new time .
            if time > 60:
                time_count += 1
                if time_count > timeTracker:
                    timeTracker = time_count

           # If left = 0 run a new try.
            if left == 0:
                left_count += 1
                if left_count > leftTracker:
                    leftTracker = left_count

            # If speed = 0 have a punish.
            if speed == 0:
                speed_count += 1
                fitness_current -= 10

                        
            if done == True:
                print(genome_id, fitness_current)

            genome.fitness = fitness_current
            


# 加载配置文件
config = neat.Config(neat.DefaultGenome, neat.DefaultReproduction,
                     neat.DefaultSpeciesSet, neat.DefaultStagnation,
                    'D:/2024/hang-on/config-feedforward')


checkpoint_file = 'neat-checkpoint-1448'    #n为第n次训练结果
p = neat.Population(config)
p = neat.Checkpointer.restore_checkpoint(checkpoint_file)

# 创建种群
# p = neat.Population(config)


# 添加报告器
p.add_reporter(neat.StdOutReporter(True))
stats = neat.StatisticsReporter()
p.add_reporter(stats)
p.add_reporter(neat.Checkpointer(10))

# 运行种群
winner = p.run(eval_genomes)

# 将优胜基因保存到文件
with open('winner.pkl', 'wb') as output:
    pickle.dump(winner, output, 1)


# checkpoint_file = 'neat-checkpoint-177'    #n为第n次训练结果
# p = neat.Population(config)
# p = neat.Checkpointer.restore_checkpoint(checkpoint_file)


env.close()